#
from typing import List, Dict, Any
from pathlib import Path
import json
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.function_registry import FunctionRegistry
from apps.cpm.ce.atom_ce import AtomCe
from apps.cpm.cpm_engine import CpmEngine
from apps.cpm.tes.range_utils import RangeUtils
from apps.cpm.tes.snr_range_utils import SnrRangeUtils
from core.radar_core import RadarCore
from common.app_registry import AppRegistry as AC
from apps.cpm.mcp.mcpc_fm import McpcFm

class SnrRangesAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.snr_ranges_agent.SnrRangesAgent'

    def startup(self, context:RlmContext) -> int:
        SnrRangesAgent.run(context=context)

    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        observes = SnrRangesAgent._observe(context=context)
        plans = SnrRangesAgent._reason(observes=observes, context=context)
        return SnrRangesAgent._act(observes=observes, context=context, plans=plans)

    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        print(f'current_input: {context.current_input}; ???????????????')
        atom_prompt = AtomCe.generate(query=context.current_input)
        print(f'{atom_prompt}')
        resp = CpmEngine.infer(query=atom_prompt)
        print(f'response:\n{resp};')
        params = json.loads(resp[8:-3])
        context.states['params'].update(params)
        observes['pt'] = params['pt']
        observes['freq'] = params['freq']
        observes['sigma'] = params['sigma']
        observes['g'] = params['g']
        observes['te'] = params['te']
        observes['b'] = params['b']
        observes['nf'] = params['nf']
        observes['loss'] = params['loss']
        observes['ranges'] = params['ranges']
        observes['range_values'] = None # RangeUtils.divide_ranges(params['ranges'][0], params['ranges'][1], params['ranges'][2])
        observes['ranges_km'] = None #observes['range_values']/1000.0
        observes['sigma_values'] = params['sigma']
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = [
            {'desc': '将距离范围划分为指定等份', 'params': {'near': observes['ranges'][0], 'far': observes['ranges'][1], 'num': observes['ranges'][2]}} ,
            {'desc': '使用雷达方程计算信噪比，在不同RCS下与距离的关系', 'params': {'context': context}}, 
            {'desc': '使用雷达方程计算信噪比，在不同发射功率下与距离的关系', 'params': {'context': context}},
            {'desc': '绘制信噪比与距离间关系', 
                'params': {
                    'x': None, #context.states['params']['ranges_km'], 
                    'ys': None #[context.states['params']['snr_pts_ranges'], context.states['params']['snr_rcss_ranges']]
                }
            }
        ]
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
        # 规划执行步骤
        # plans = self.plan_processing(observes)
        for plan in plans:
            print(f'### {plan["desc"]}... v0.0.2')
            fc_prompt = f'{FunctionRegistry.FC_PROMPT}\n{plan["desc"]}'
            func_name_raw = CpmEngine.infer(query=fc_prompt)
            print(f'@@@ func_name_raw: \n{func_name_raw};')
            func_name_aux = func_name_raw.split('<|func_name|_start>')[-1]
            func_name = func_name_aux.split('<|func_name|_end>')[0]
            for key in plan['params'].keys():
                if key in context.states['params']:
                    print(f'@@@ {key}: ')
                    plan['params'][key] = context.states['params'][key]
            if FunctionRegistry.FUNCTIONS[func_name]['type'] == 'function':
                rsts = FunctionRegistry.FUNCTIONS[func_name]['function'](**plan['params'])
            elif FunctionRegistry.FUNCTIONS[func_name]['type'] == 'agent':
                rsts = FunctionRegistry.FUNCTIONS[func_name]['function'](**plan['params'])
            elif FunctionRegistry.FUNCTIONS[func_name]['type'] == 'mcp':
                rsts = FunctionRegistry.FUNCTIONS[func_name]['function'](**plan['params'])
            if 'params' in rsts:
                observes.update(rsts['params'])
                context.states['params'].update(observes)

        # # fig, ax, axs = SnrRangeUtils.init_matplotlib(rows=2, cols=1)
        # # generate_range 功能调用
        # # snr_range_rcs
        # # 内部Agent调用
        # # snr_matrix1 = np.array([RadarCore.radar_eq(observes['pt'][0], observes['freq'], observes['g'], s, observes['te'], observes['b'], observes['nf'], observes['loss'], observes['ranges_km']) 
        # #                     for s in observes['sigma_values']])
        # snr_matrix1 = context.states['params']['snr_rcss_ranges']
        # # MCP服务器调用
        # SnrRangeUtils.create_radar_plot(AC.axs[0], observes['ranges_km'], snr_matrix1, title='信号噪比距离关系（不同RCS）', labels=[r'$\sigma$ = 0 dBsm', 
        #         r'$\sigma$ = -10 dBsm', 
        #         r'$\sigma$ = +10 dBsm'])
        # # snr_range_pt
        # # snr_matrix2 = np.array([RadarCore.radar_eq(p, observes['freq'], observes['g'], observes['sigma_values'][0], observes['te'], observes['b'], observes['nf'], observes['loss'], observes['ranges_km']) 
        # #                     for p in observes['pt']])
        # snr_matrix2 = context.states['params']['snr_pts_ranges']
        # # 绘制第二个图（不同功率）
        # SnrRangeUtils.create_radar_plot(
        #     AC.axs[1], observes['ranges_km'], snr_matrix2,
        #     title='信噪比与距离关系图（不同发射功率）',
        #     labels=['Pt = 1.5 MW', 
        #         'Pt = 0.6 MW', 
        #         'Pt = 2.7 MW']
        # )
        # # 统一调整
        # plt.tight_layout()
        # plt.show()
        context.states['params'].update(observes)
        print(f'最终图片位置：img_fn: {context.states["params"]["img_fn"]}')
        # 创建目录: work/reports/v001
        base_folder = 'work/reports/v001'
        Path(base_folder).mkdir(parents=True, exist_ok=True)
        # 创建图片目录：work/reports/v001/images
        images_folder = f'{base_folder}/images'
        Path(images_folder).mkdir(parents=True, exist_ok=True)
        # 下载图片文件：img_fn
        content = McpcFm.download_file(context.states["params"]["img_fn"])
        with open(f'{images_folder}/img_001.png', 'wb') as wfd:
            wfd.write(content)
        # 创建md文件./work/reports/v001/report.md，将user_query写入md文件，将图像写入md文件
        with open(f'{base_folder}/report.md', 'w', encoding='utf-8') as wfd:
            wfd.write(f'# 雷达设计：信噪比与距离实验报告\n')
            wfd.write(context.current_input)
            wfd.write(f'\n![img_001.png](./images/img_001.png "")')
        return {'result': 0, 'params': {}}